Faster model inference on CPU with Neuralmagic
An example using YOLOv5 and Neuralmagic.
• Jaco Verster • 18 min read
- Step 1: Download a dataset
- Step 2: Explore the dataset
- Step 3: Split the dataset into train/test/val subsets
- Step 4: Install DeepSparse and SparseML libraries and train our model
- Step 4: Export the Sparse Model to ONNX format for CPU inference
- Step 5: DeepSparse Engine model deployment
- Ste 6: Run the model on a youtube video using the Ultralytics YOLO library
Step 1: Download a dataset
Download the Guns dataset from Roboflow's public databse.
Select YOLOv5 from the download formats and copy your unique terminal command into the cell below.
! curl -L "https://public.roboflow.com/ds/xXg040ThIk?key=mM2BVJ4kkZ" > data.zip \
&& unzip -q data.zip -d downloads \
&& rm data.zip
% Total % Received % Xferd Average Speed Time Time Time Current
Dload Upload Total Spent Left Speed
100 901 100 901 0 0 3575 0 --:--:-- --:--:-- --:--:-- 3575
100 58.4M 100 58.4M 0 0 70.8M 0 --:--:-- --:--:-- --:--:-- 70.8M
! pip install -q opencv-python-headless==4.5.4.60 fiftyone
import fiftyone as fo
data_path = "downloads/export/images"
labels_path = "downloads/export/labels"
classes = ["pistol"]
# Import dataset by explicitly providing paths to the source media and labels
dataset = fo.Dataset.from_dir(
dataset_type=fo.types.YOLOv4Dataset,
data_path=data_path,
label_field="ground_truth",
labels_path=labels_path,
classes=classes,
)
# View data
session = fo.launch_app(dataset)
import fiftyone.utils.splits as fous
import os
def split_dataset(dataset, split_dict, output_path):
fous.random_split(dataset, split_dict, seed=42)
for tag in list(split_dict.keys()):
view = dataset.match_tags(tag)
# Export each split separately
dir = os.path.join(output_path, tag)
# Export the dataset
view.export(
export_dir=dir,
dataset_type=fo.types.YOLOv4Dataset,
label_field="ground_truth",
)
split_dataset(
dataset=dataset,
split_dict={"train": 0.8, "test": 0.1, "val": 0.1},
output_path="dataset/guns/"
)
100% |███████████████| 2377/2377 [2.7s elapsed, 0s remaining, 859.1 samples/s] 100% |█████████████████| 297/297 [355.2ms elapsed, 0s remaining, 836.1 samples/s] 100% |█████████████████| 297/297 [397.3ms elapsed, 0s remaining, 747.5 samples/s]
content = """# Train/val/test sets
path: /content/dataset/guns # dataset root dir
train: train/data # train images (relative to 'path')
test: test/data # test images (relative to 'path')
val: val/data # val images (relative to 'path')
# classes
nc: 1 # number of classes
names: ['pistol']"""
with open("pistols.yml", "w") as f:
f.write(content)
! rm -r downloads
! pip install -q deepsparse[yolo] sparseml[torchvision]
Grab the transfer learning reciepe from github
! curl -LO "https://raw.githubusercontent.com/neuralmagic/sparseml/main/integrations/ultralytics-yolov5/recipes/yolov5.transfer_learn_pruned.md"
% Total % Received % Xferd Average Speed Time Time Time Current
Dload Upload Total Spent Left Speed
100 3084 100 3084 0 0 16146 0 --:--:-- --:--:-- --:--:-- 16146
Train the model using transfer learning on yolov5s (small) and the downloaded reciepe.
The batch-size = -1 parameter turns on autobatch.
! sparseml.yolov5.train \
--data pistols.yml \
--cfg models_v5.0/yolov5s.yaml \
--weights zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96 \
--hyp data/hyps/hyp.finetune.yaml \
--batch-size -1 \
--imgsz 416 \
--recipe yolov5.transfer_learn_pruned.md
train: weights=zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned-aggressive_96, cfg=models_v5.0/yolov5s.yaml, data=pistols.yml, hyp=data/hyps/hyp.finetune.yaml, epochs=300, batch_size=-1, imgsz=416, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, evolve=None, bucket=, cache=None, image_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=yolov5_runs/train, name=exp, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=0, freeze=[0], save_period=-1, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest, recipe=yolov5.transfer_learn_pruned.md, disable_ema=False, max_train_steps=-1, max_eval_steps=-1, one_shot=False, num_export_samples=0 github: skipping check (not a git repository), for updates see https://github.com/ultralytics/yolov5 fatal: not a git repository (or any of the parent directories): .git YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CUDA:0 (Tesla T4, 15110MiB) hyperparameters: lr0=0.0032, lrf=0.12, momentum=0.843, weight_decay=0.00036, warmup_epochs=2.0, warmup_momentum=0.5, warmup_bias_lr=0.05, box=0.0296, cls=0.243, cls_pw=0.631, obj=0.301, obj_pw=0.911, iou_t=0.2, anchor_t=2.91, fl_gamma=0.0, hsv_h=0.0138, hsv_s=0.664, hsv_v=0.464, degrees=0.373, translate=0.245, scale=0.898, shear=0.602, perspective=0.0, flipud=0.00856, fliplr=0.5, mosaic=1.0, mixup=0.243, copy_paste=0.0 Weights & Biases: run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED) TensorBoard: Start with 'tensorboard --logdir yolov5_runs/train', view at http://localhost:6006/ Obtaining new sparse zoo credentials token Getting signed url for c13e55cb-dd6c-4492-a079-8986af0b65e6/model.pt Downloading model file model.pt to /root/.cache/sparsezoo/c13e55cb-dd6c-4492-a079-8986af0b65e6/pytorch/model.pt downloading...: 100% 14.1M/14.1M [00:00<00:00, 37.0MB/s] Overriding model.yaml nc=80 with nc=1 from n params module arguments 0 -1 1 3520 yolov5.models.common.Focus [3, 32, 3] 1 -1 1 18560 yolov5.models.common.Conv [32, 64, 3, 2] 2 -1 1 18816 yolov5.models.common.C3 [64, 64, 1] 3 -1 1 73984 yolov5.models.common.Conv [64, 128, 3, 2] 4 -1 3 156928 yolov5.models.common.C3 [128, 128, 3] 5 -1 1 295424 yolov5.models.common.Conv [128, 256, 3, 2] 6 -1 3 625152 yolov5.models.common.C3 [256, 256, 3] 7 -1 1 1180672 yolov5.models.common.Conv [256, 512, 3, 2] 8 -1 1 656896 yolov5.models.common.SPP [512, 512, [5, 9, 13]] 9 -1 1 1182720 yolov5.models.common.C3 [512, 512, 1, False] 10 -1 1 131584 yolov5.models.common.Conv [512, 256, 1, 1] 11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest'] 12 [-1, 6] 1 0 yolov5.models.common.Concat [1] 13 -1 1 361984 yolov5.models.common.C3 [512, 256, 1, False] 14 -1 1 33024 yolov5.models.common.Conv [256, 128, 1, 1] 15 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest'] 16 [-1, 4] 1 0 yolov5.models.common.Concat [1] 17 -1 1 90880 yolov5.models.common.C3 [256, 128, 1, False] 18 -1 1 147712 yolov5.models.common.Conv [128, 128, 3, 2] 19 [-1, 14] 1 0 yolov5.models.common.Concat [1] 20 -1 1 296448 yolov5.models.common.C3 [256, 256, 1, False] 21 -1 1 590336 yolov5.models.common.Conv [256, 256, 3, 2] 22 [-1, 10] 1 0 yolov5.models.common.Concat [1] 23 -1 1 1182720 yolov5.models.common.C3 [512, 512, 1, False] 24 [17, 20, 23] 1 16182 yolov5.models.yolo.Detect [1, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]] overriding activations in model to Hardswish YOLOv5s summary: 283 layers, 7063542 parameters, 7063542 gradients, 16.5 GFLOPs Transferred 354/361 items from /root/.cache/sparsezoo/c13e55cb-dd6c-4492-a079-8986af0b65e6/pytorch/model.pt AutoBatch: Computing optimal batch size for --imgsz 416 AutoBatch: CUDA:0 (Tesla T4) 14.76G total, 0.06G reserved, 0.05G allocated, 14.64G free Params GFLOPs GPU_mem (GB) forward (ms) backward (ms) input output 7063542 6.961 0.201 16.78 14.15 (1, 3, 416, 416) list 7063542 13.92 0.310 17.35 15.45 (2, 3, 416, 416) list 7063542 27.85 0.581 18.58 18.79 (4, 3, 416, 416) list 7063542 55.69 1.126 19.02 26.23 (8, 3, 416, 416) list 7063542 111.4 2.223 32.87 50.45 (16, 3, 416, 416) list AutoBatch: Using batch-size 96 for CUDA:0 13.28G/14.76G (90%) Scaled weight_decay = 0.00054 optimizer: SGD with parameter groups 59 weight (no decay), 62 weight, 62 bias albumentations: version 1.0.3 required by YOLOv5, but version 0.1.12 is currently installed train: Scanning '/content/dataset/guns/train/data' images and labels...2377 found, 0 missing, 0 empty, 0 corrupt: 100% 2377/2377 [00:01<00:00, 1379.73it/s] train: New cache created: /content/dataset/guns/train/data.cache val: Scanning '/content/dataset/guns/val/data' images and labels...297 found, 0 missing, 0 empty, 0 corrupt: 100% 297/297 [00:00<00:00, 634.83it/s] val: New cache created: /content/dataset/guns/val/data.cache Plotting labels to yolov5_runs/train/exp2/labels.jpg... AutoAnchor: 2.95 anchors/target, 1.000 Best Possible Recall (BPR). Current anchors are a good fit to dataset ✅ Image sizes 416 train, 416 val Using 2 dataloader workers Logging results to yolov5_runs/train/exp2 Starting training for 300 epochs... Disabling LR scheduler, managing LR using SparseML recipe Overriding number of epochs from SparseML manager to 50 Epoch gpu_mem box obj cls labels img_size 0/49 12.2G 0.07143 0.005311 0 243 416: 100% 25/25 [00:35<00:00, 1.43s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:04<00:00, 2.16s/it] all 297 357 0.0179 0.185 0.0106 0.00189 Epoch gpu_mem box obj cls labels img_size 1/49 14.3G 0.0591 0.005631 0 237 416: 100% 25/25 [00:32<00:00, 1.30s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.95s/it] all 297 357 0.11 0.199 0.0727 0.0206 Epoch gpu_mem box obj cls labels img_size 2/49 12.3G 0.0511 0.005889 0 186 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.61s/it] all 297 357 0.356 0.401 0.275 0.0831 Epoch gpu_mem box obj cls labels img_size 3/49 12.9G 0.04455 0.006083 0 195 416: 100% 25/25 [00:33<00:00, 1.33s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.53s/it] all 297 357 0.419 0.473 0.406 0.132 Epoch gpu_mem box obj cls labels img_size 4/49 12.9G 0.03806 0.006098 0 233 416: 100% 25/25 [00:32<00:00, 1.32s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.50s/it] all 297 357 0.555 0.605 0.598 0.259 Epoch gpu_mem box obj cls labels img_size 5/49 12.9G 0.03427 0.005962 0 258 416: 100% 25/25 [00:32<00:00, 1.30s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.46s/it] all 297 357 0.729 0.591 0.686 0.347 Epoch gpu_mem box obj cls labels img_size 6/49 12.9G 0.03197 0.00576 0 207 416: 100% 25/25 [00:33<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.47s/it] all 297 357 0.718 0.669 0.718 0.392 Epoch gpu_mem box obj cls labels img_size 7/49 12.9G 0.03004 0.005625 0 254 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.52s/it] all 297 357 0.749 0.669 0.737 0.411 Epoch gpu_mem box obj cls labels img_size 8/49 13G 0.02922 0.005537 0 235 416: 100% 25/25 [00:34<00:00, 1.38s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.44s/it] all 297 357 0.701 0.691 0.752 0.436 Epoch gpu_mem box obj cls labels img_size 9/49 13G 0.02814 0.005131 0 196 416: 100% 25/25 [00:33<00:00, 1.35s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.47s/it] all 297 357 0.717 0.709 0.772 0.447 Epoch gpu_mem box obj cls labels img_size 10/49 13G 0.02756 0.005152 0 247 416: 100% 25/25 [00:34<00:00, 1.37s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.48s/it] all 297 357 0.671 0.787 0.78 0.461 Epoch gpu_mem box obj cls labels img_size 11/49 13G 0.02719 0.005132 0 227 416: 100% 25/25 [00:32<00:00, 1.30s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.51s/it] all 297 357 0.847 0.666 0.8 0.463 Epoch gpu_mem box obj cls labels img_size 12/49 13G 0.02633 0.004897 0 265 416: 100% 25/25 [00:32<00:00, 1.32s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.48s/it] all 297 357 0.722 0.773 0.811 0.495 Epoch gpu_mem box obj cls labels img_size 13/49 13G 0.02559 0.005043 0 196 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.41s/it] all 297 357 0.691 0.814 0.791 0.469 Epoch gpu_mem box obj cls labels img_size 14/49 13G 0.02537 0.004815 0 232 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.44s/it] all 297 357 0.791 0.731 0.811 0.499 Epoch gpu_mem box obj cls labels img_size 15/49 13G 0.02512 0.004849 0 230 416: 100% 25/25 [00:34<00:00, 1.37s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.45s/it] all 297 357 0.765 0.739 0.81 0.483 Epoch gpu_mem box obj cls labels img_size 16/49 13G 0.02504 0.004759 0 206 416: 100% 25/25 [00:33<00:00, 1.33s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.47s/it] all 297 357 0.737 0.779 0.82 0.512 Epoch gpu_mem box obj cls labels img_size 17/49 13G 0.02469 0.004673 0 267 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.41s/it] all 297 357 0.802 0.72 0.822 0.512 Epoch gpu_mem box obj cls labels img_size 18/49 13G 0.02447 0.004662 0 231 416: 100% 25/25 [00:33<00:00, 1.35s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.46s/it] all 297 357 0.778 0.759 0.824 0.522 Epoch gpu_mem box obj cls labels img_size 19/49 13G 0.02424 0.004754 0 235 416: 100% 25/25 [00:33<00:00, 1.35s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.39s/it] all 297 357 0.785 0.759 0.827 0.533 Epoch gpu_mem box obj cls labels img_size 20/49 13G 0.02374 0.004643 0 213 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.42s/it] all 297 357 0.857 0.695 0.823 0.517 Epoch gpu_mem box obj cls labels img_size 21/49 13G 0.02402 0.004533 0 209 416: 100% 25/25 [00:34<00:00, 1.37s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.46s/it] all 297 357 0.796 0.754 0.816 0.523 Epoch gpu_mem box obj cls labels img_size 22/49 13.6G 0.0235 0.004528 0 236 416: 100% 25/25 [00:32<00:00, 1.31s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.40s/it] all 297 357 0.777 0.782 0.833 0.545 Epoch gpu_mem box obj cls labels img_size 23/49 13.6G 0.02366 0.004616 0 264 416: 100% 25/25 [00:33<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.44s/it] all 297 357 0.816 0.748 0.832 0.531 Epoch gpu_mem box obj cls labels img_size 24/49 13.6G 0.02322 0.004491 0 238 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.49s/it] all 297 357 0.814 0.748 0.829 0.519 Epoch gpu_mem box obj cls labels img_size 25/49 13.6G 0.02305 0.004363 0 213 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.40s/it] all 297 357 0.768 0.773 0.835 0.536 Epoch gpu_mem box obj cls labels img_size 26/49 13.6G 0.02326 0.004525 0 238 416: 100% 25/25 [00:33<00:00, 1.33s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.43s/it] all 297 357 0.801 0.787 0.849 0.552 Epoch gpu_mem box obj cls labels img_size 27/49 13.6G 0.02263 0.004288 0 250 416: 100% 25/25 [00:33<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.38s/it] all 297 357 0.813 0.77 0.845 0.548 Epoch gpu_mem box obj cls labels img_size 28/49 13.6G 0.02294 0.004502 0 245 416: 100% 25/25 [00:33<00:00, 1.33s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.46s/it] all 297 357 0.821 0.762 0.84 0.546 Epoch gpu_mem box obj cls labels img_size 29/49 13.6G 0.02249 0.004358 0 223 416: 100% 25/25 [00:33<00:00, 1.35s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.42s/it] all 297 357 0.811 0.78 0.839 0.547 Epoch gpu_mem box obj cls labels img_size 30/49 13.7G 0.0226 0.004365 0 234 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.40s/it] all 297 357 0.832 0.75 0.837 0.545 Epoch gpu_mem box obj cls labels img_size 31/49 13.7G 0.02216 0.004202 0 218 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.39s/it] all 297 357 0.847 0.744 0.841 0.542 Epoch gpu_mem box obj cls labels img_size 32/49 13.7G 0.02241 0.004421 0 258 416: 100% 25/25 [00:33<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.45s/it] all 297 357 0.789 0.798 0.841 0.553 Epoch gpu_mem box obj cls labels img_size 33/49 13.7G 0.02187 0.004324 0 183 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.42s/it] all 297 357 0.798 0.798 0.847 0.551 Epoch gpu_mem box obj cls labels img_size 34/49 13.7G 0.02235 0.004243 0 217 416: 100% 25/25 [00:32<00:00, 1.30s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.45s/it] all 297 357 0.821 0.773 0.843 0.55 Epoch gpu_mem box obj cls labels img_size 35/49 13.7G 0.02182 0.004277 0 234 416: 100% 25/25 [00:33<00:00, 1.35s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.42s/it] all 297 357 0.804 0.796 0.849 0.559 Epoch gpu_mem box obj cls labels img_size 36/49 13.7G 0.02203 0.004312 0 260 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.47s/it] all 297 357 0.809 0.784 0.846 0.553 Epoch gpu_mem box obj cls labels img_size 37/49 13.7G 0.02222 0.00428 0 200 416: 100% 25/25 [00:32<00:00, 1.31s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.43s/it] all 297 357 0.806 0.804 0.85 0.557 Epoch gpu_mem box obj cls labels img_size 38/49 13.7G 0.02202 0.004223 0 210 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.46s/it] all 297 357 0.784 0.796 0.836 0.551 Epoch gpu_mem box obj cls labels img_size 39/49 13.7G 0.02193 0.004199 0 256 416: 100% 25/25 [00:34<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.43s/it] all 297 357 0.804 0.795 0.844 0.556 Epoch gpu_mem box obj cls labels img_size 40/49 13.7G 0.02192 0.004291 0 246 416: 100% 25/25 [00:34<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.24s/it] all 297 357 0.792 0.8 0.839 0.558 Epoch gpu_mem box obj cls labels img_size 41/49 13.7G 0.02164 0.004269 0 235 416: 100% 25/25 [00:32<00:00, 1.31s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.47s/it] all 297 357 0.783 0.81 0.842 0.561 Epoch gpu_mem box obj cls labels img_size 42/49 13.7G 0.02207 0.004228 0 247 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.43s/it] all 297 357 0.79 0.812 0.845 0.562 Epoch gpu_mem box obj cls labels img_size 43/49 13.7G 0.02193 0.004293 0 263 416: 100% 25/25 [00:33<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.57s/it] all 297 357 0.796 0.801 0.85 0.561 Epoch gpu_mem box obj cls labels img_size 44/49 13.7G 0.02202 0.004254 0 242 416: 100% 25/25 [00:34<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.44s/it] all 297 357 0.81 0.793 0.849 0.556 Epoch gpu_mem box obj cls labels img_size 45/49 13.7G 0.02163 0.004189 0 209 416: 100% 25/25 [00:33<00:00, 1.33s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.48s/it] all 297 357 0.791 0.807 0.849 0.557 Epoch gpu_mem box obj cls labels img_size 46/49 13.7G 0.02168 0.004311 0 241 416: 100% 25/25 [00:34<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.44s/it] all 297 357 0.801 0.801 0.844 0.556 Epoch gpu_mem box obj cls labels img_size 47/49 13.7G 0.02202 0.004218 0 231 416: 100% 25/25 [00:33<00:00, 1.36s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.44s/it] all 297 357 0.811 0.793 0.847 0.559 Epoch gpu_mem box obj cls labels img_size 48/49 13.7G 0.02156 0.004261 0 230 416: 100% 25/25 [00:32<00:00, 1.32s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.51s/it] all 297 357 0.835 0.765 0.842 0.556 Epoch gpu_mem box obj cls labels img_size 49/49 13.7G 0.02169 0.004359 0 239 416: 100% 25/25 [00:33<00:00, 1.34s/it] Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:02<00:00, 1.43s/it] all 297 357 0.803 0.81 0.851 0.561 51 epochs completed in 0.526 hours. Optimizer stripped from yolov5_runs/train/exp2/weights/last.pt, 14.4MB Optimizer stripped from yolov5_runs/train/exp2/weights/best.pt, 14.4MB Validating yolov5_runs/train/exp2/weights/best.pt... Fusing layers... YOLOv5s summary: 224 layers, 7053910 parameters, 0 gradients, 16.3 GFLOPs Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 2/2 [00:03<00:00, 1.97s/it] all 297 357 0.789 0.815 0.845 0.562 Results saved to yolov5_runs/train/exp2
! sparseml.yolov5.export_onnx \
--weights yolov5_runs/train/exp/weights/best.pt \
--imgsz 416 \
--dynamic
export: data=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/data/coco128.yaml, weights=['yolov5_runs/train/exp/weights/best.pt'], imgsz=[416], batch_size=1, device=cpu, half=False, inplace=False, train=False, optimize=False, int8=False, dynamic=True, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, remove_grid=False, include=['torchscript', 'onnx'] export: data=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/data/coco128.yaml, weights=['yolov5_runs/train/exp/weights/best.pt'], imgsz=[416], batch_size=1, device=cpu, half=False, inplace=False, train=False, optimize=False, int8=False, dynamic=True, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, remove_grid=False, include=['torchscript', 'onnx'] fatal: not a git repository (or any of the parent directories): .git YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CPU Fusing layers... YOLOv5s summary: 224 layers, 7053910 parameters, 0 gradients, 16.3 GFLOPs PyTorch: starting from yolov5_runs/train/exp/weights/best.pt with output shape (1, 10647, 6) (13.8 MB) TorchScript: starting export with torch 1.9.1+cu102... TorchScript: export success, saved as yolov5_runs/train/exp/weights/best.torchscript (27.2 MB) ONNX: starting export with onnx 1.10.1... ONNX: export success, saved as yolov5_runs/train/exp/weights/best.onnx (27.0 MB) Export complete (7.26s) Results saved to /content/yolov5_runs/train/exp/weights Detect: python detect.py --weights yolov5_runs/train/exp/weights/best.onnx PyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5_runs/train/exp/weights/best.onnx') Validate: python val.py --weights yolov5_runs/train/exp/weights/best.onnx Visualize: https://netron.app
! sparseml.yolov5.val_onnx \
--model_path yolov5_runs/train/exp/weights/best.onnx \
--batch-size 33 \
--imgsz 416 \
--task test \
--data pistols.yml
val_onnx: data=pistols.yml, model_path=yolov5_runs/train/exp/weights/best.onnx, batch_size=33, imgsz=416, conf_thres=0.001, iou_thres=0.6, task=test, device=, workers=8, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val, name=exp, exist_ok=False, half=False, dnn=False, engine=deepsparse, num_cores=None val_onnx: data=pistols.yml, model_path=yolov5_runs/train/exp/weights/best.onnx, batch_size=33, imgsz=416, conf_thres=0.001, iou_thres=0.6, task=test, device=, workers=8, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=False, project=../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val, name=exp, exist_ok=False, half=False, dnn=False, engine=deepsparse, num_cores=None fatal: not a git repository (or any of the parent directories): .git YOLOv5 🚀 2022-6-30 torch 1.9.1+cu102 CUDA:0 (Tesla T4, 15110MiB) DeepSparse Engine, Copyright 2021-present / Neuralmagic, Inc. version: 1.0.0 (8eaddc24) (release) (optimized) (system=avx2, binary=avx2) test: Scanning '/content/dataset/guns/test/data' images and labels...297 found, 0 missing, 0 empty, 0 corrupt: 100% 297/297 [00:00<00:00, 1260.64it/s] test: New cache created: /content/dataset/guns/test/data.cache Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 9/9 [00:34<00:00, 3.88s/it] /usr/local/lib/python3.7/dist-packages/yolov5/utils/metrics.py:74: FutureWarning: elementwise comparison failed; returning scalar instead, but in the future will perform elementwise comparison names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data all 297 351 0.918 0.764 0.886 0.578 Speed: 0.0ms pre-process, 116.0ms inference, 0.0ms NMS per image at shape (33, 3, 416, 416) Results saved to ../usr/local/lib/python3.7/dist-packages/sparseml/yolov5/runs/val/exp2
Ste 6: Run the model on a youtube video using the Ultralytics YOLO library
We can also run the model trained with Neuralmagic CLI in using the Ultralytics YOLO library. Simply install the library and run inference using the *.pt model weights.
For this example we will use a scene from John Wick making Revolver to test the model.
! git clone https://github.com/ultralytics/yolov5
! pip install -q -r yolov5/requirements.txt
Cloning into 'yolov5'... remote: Enumerating objects: 12351, done. remote: Counting objects: 100% (16/16), done. remote: Compressing objects: 100% (15/15), done. remote: Total 12351 (delta 6), reused 4 (delta 1), pack-reused 12335 Receiving objects: 100% (12351/12351), 12.18 MiB | 31.10 MiB/s, done. Resolving deltas: 100% (8520/8520), done.
! pip install -q -r yolov5/requirements.txt
!python yolov5/detect.py \
--weights /content/yolov5_runs/train/exp/weights/best.pt \
--source "https://www.youtube.com/watch?v=nB0A4pFIOA0" \
--conf-thres=0.6 \
--imgsz 416
Below are a few screenshots from the annotated video.
